package com.yuluo.aihelper.ai.rag;


import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import jakarta.annotation.Resource;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.ArrayList;
import java.util.List;

/**
 * 加载 RAG
 */
@Configuration
public class RagConfig {
    @Resource
    private EmbeddingModel qwenEmbeddingModel;

    @Resource
    private EmbeddingStore<TextSegment> embeddingStore;

    @Bean
    public ContentRetriever contentRetriever(){
        // 加载 sow 文件夹下的文档
        List<Document> sowDocuments = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs/sow");
        // 加载 maid 文件夹下的文档
        List<Document> maidDocuments = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs/maid");

        // 合并文档列表
        List<Document> allDocuments = new ArrayList<>();
        allDocuments.addAll(sowDocuments);
        allDocuments.addAll(maidDocuments);

        //2.文档分割：基于段落分割，按照最大长度1000字符，重叠部分最大200字符切分
        DocumentByParagraphSplitter paragraphSplitter = new DocumentByParagraphSplitter(1000, 200);
        //3.自定义文档加载器
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(paragraphSplitter)
                //在原始文本前加上文件名作为上下文信息，有助于后续检索时识别来源。结果格式：文件名\n + 段落文本
                .textSegmentTransformer(textSegment -> TextSegment.from(
                        textSegment.metadata().getString("file_name") + "\n" + textSegment.text(),
                        textSegment.metadata()))
                .embeddingModel(qwenEmbeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        //数据摄入操作
        ingestor.ingest(allDocuments);
        //4.自定义内容查询器
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .maxResults(5)//每次检索最多返回5个结果
                .minScore(0.75)
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)//最低相似度得分阈值，低于该分数的结果会被过滤
                .build();
        return contentRetriever;
    }
}






















